import os
import numpy as np
import cv2
import argparse
import pyclipper
import bpu_infer_lib
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt
import collections


class DetectionModel:
    def __init__(self, model_path, threshold=0.5, ratio_prime=2.7, input_size=(640, 640)):
        """
        Initialize the detection model.
        """
        try:
            self.model = bpu_infer_lib.Infer(False)
            if not self.model.load_model(model_path):
                raise RuntimeError(f"Failed to load model from '{model_path}'.")
        except Exception as e:
            print(f"初始化检测模型时出错: {e}")
            raise
        
        self.threshold = threshold
        self.ratio_prime = ratio_prime
        self.input_size = input_size

    def predict(self, img, img_path, min_area=100):
        """
        Predict bounding boxes for the given image.
        """
        try:
            img_shape = img.shape[:2]
            
            self.model.read_img_to_nv12(img_path, 0)
            self.model.forward()
            
            # 获取推理结果
            preds = self.model.get_infer_res_np_float32(0).reshape(1, *self.input_size)
            preds = np.where(preds > self.threshold, 255, 0).astype(np.uint8).squeeze()
            preds = cv2.resize(preds, (img_shape[1], img_shape[0]))

            contours, _ = cv2.findContours(preds, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            dilated_polys = self._dilate_contours(contours)
            boxes_list = self._get_bounding_boxes(dilated_polys, min_area)

            return dilated_polys, boxes_list
        except Exception as e:
            print(f"检测预测时出错: {e}")
            return [], []

    def _dilate_contours(self, contours):
        """Dilate contours using the ratio_prime."""
        dilated_polys = []
        for poly in contours:
            try:
                poly = poly[:, 0, :]
                arc_length = cv2.arcLength(poly, True)
                if arc_length == 0:
                    continue
                D_prime = (cv2.contourArea(poly) * self.ratio_prime / arc_length)

                pco = pyclipper.PyclipperOffset()
                pco.AddPath(poly, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
                dilated_poly = np.array(pco.Execute(D_prime))

                if dilated_poly.size == 0 or dilated_poly.dtype != np.int_ or len(dilated_poly) != 1:
                    continue
                dilated_polys.append(dilated_poly)
            except Exception as e:
                print(f"轮廓膨胀时出错: {e}")
                continue
        return dilated_polys

    def _get_bounding_boxes(self, dilated_polys, min_area):
        """Get bounding boxes from dilated polygons."""
        boxes_list = []
        for cnt in dilated_polys:
            try:
                if cv2.contourArea(cnt) < min_area:
                    continue
                rect = cv2.minAreaRect(cnt)
                box = cv2.boxPoints(rect).astype(np.int_)
                boxes_list.append(box)
            except Exception as e:
                print(f"获取边界框时出错: {e}")
                continue
        return boxes_list


class strLabelConverter:
    """Convert between string and label for OCR tasks."""

    def __init__(self, alphabet, ignore_case=False):
        self._ignore_case = ignore_case
        if self._ignore_case:
            alphabet = alphabet.lower()
        self.alphabet = alphabet + '-'  # Add a special '-' character for padding

        # Create a dictionary mapping characters to indices
        self.dict = {}
        for i, char in enumerate(alphabet):
            # Note: 0 is reserved for 'blank' required by CTC loss
            self.dict[char] = i + 1

    def encode(self, text):
        """Encode a string or a list of strings into a sequence of indices."""
        if isinstance(text, str):
            text = [
                self.dict[char.lower() if self._ignore_case else char]
                for char in text
            ]
            length = [len(text)]
        elif isinstance(text, collections.Iterable):
            length = [len(s) for s in text]
            text = ''.join(text)
            text, _ = self.encode(text)
        return np.array(text, dtype=np.int32), np.array(length, dtype=np.int32)

    def decode(self, t, length, raw=False):
        """Decode a sequence of indices back into a string."""
        if len(length) == 1:
            length = length[0]
            assert len(t) == length, f"text with length: {len(t)} does not match declared length: {length}"
            if raw:
                return ''.join([self.alphabet[i - 1] for i in t if i > 0])
            else:
                char_list = []
                for i in range(length):
                    if t[i] != 0 and (not (i > 0 and t[i - 1] == t[i])):
                        char_list.append(self.alphabet[t[i] - 1])
                return ''.join(char_list)
        else:
            # Batch decoding mode
            assert len(t) == length.sum(), f"texts with length: {len(t)} do not match declared length: {length.sum()}"
            texts = []
            index = 0
            for i in range(length.size):
                l = length[i]
                texts.append(
                    self.decode(
                        t[index:index + l], np.array([l]), raw=raw))
                index += l
            return texts


class rec_model:
    def __init__(self, model_path, converter, input_size=(48, 320), output_size=(40, 97)):
        """
        Initialize the recognition model.
        """
        try:
            self.model = bpu_infer_lib.Infer(False)
            if not self.model.load_model(model_path):
                raise RuntimeError(f"Failed to load model from '{model_path}'.")
        except Exception as e:
            print(f"初始化识别模型时出错: {e}")
            raise

        self.converter = converter
        self.output_size = output_size
        self.input_size = input_size

    def predict_float(self, img):
        """
        Perform prediction on an image using float32 input.
        """
        try:
            print(f"输入图像尺寸: {img.shape}")
            
            # 预处理图像
            if len(img.shape) == 3:
                # 如果是彩色图像，转换为灰度图
                if img.shape[2] == 3:
                    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                else:
                    img_gray = img[:, :, 0]
            else:
                img_gray = img
            
            # 调整图像尺寸
            image_resized = cv2.resize(img_gray, dsize=(self.input_size[1], self.input_size[0]))
            print(f"调整后尺寸: {image_resized.shape}")
            
            # 转换为3通道
            if len(image_resized.shape) == 2:
                image_resized = cv2.cvtColor(image_resized, cv2.COLOR_GRAY2RGB)
            
            # 归一化
            image_resized = (image_resized / 255.0).astype(np.float32)
            
            # BGR to RGB (如果需要)
            if image_resized.shape[2] == 3:
                image_resized = image_resized[:, :, [2, 1, 0]]
            
            # 转换为模型输入格式 NCHW
            input_image = image_resized[None].transpose(0, 3, 1, 2)
            print(f"模型输入尺寸: {input_image.shape}")
            
            # 保存调试图像
            debug_img = (image_resized * 255).astype(np.uint8)
            cv2.imwrite(f'debug_crop_{np.random.randint(1000)}.jpg', debug_img)
            
            # 模型推理
            self.model.read_numpy_arr_float32(input_image, 0)
            self.model.forward(more=True)
            
            # 获取推理结果
            preds = self.model.get_infer_res_np_float32(0).reshape(1, *self.output_size)
            print("预测形状:", preds.shape)
            print("预测值范围:", preds.min(), "到", preds.max())
            
            # 转置和argmax
            preds = np.transpose(preds, (1, 0, 2))  # (40, 1, 97)
            print("转置后形状:", preds.shape)
            
            # 查看预测概率分布
            max_probs = np.max(preds, axis=2)
            print("最大概率分布:", max_probs.flatten()[:10])  # 显示前10个
            
            preds = np.argmax(preds, axis=2)  # (40, 1)
            preds = preds.transpose(1, 0).reshape(-1)  # (40,)
            
            print("预测索引:", preds[:20])  # 显示前20个索引
            
            preds_size = np.array([preds.size], dtype=np.int32)
            raw_pred = self.converter.decode(np.array(preds), np.array(preds_size), raw=True)
            sim_pred = self.converter.decode(np.array(preds), np.array(preds_size), raw=False)
            
            return raw_pred, sim_pred
            
        except Exception as e:
            print(f"浮点预测时出错: {e}")
            import traceback
            traceback.print_exc()
            return "", ""


def load_image(img_path):
    """Load an image from a file path."""
    try:
        img = cv2.imread(img_path)
        if img is None:
            raise FileNotFoundError(f"图像文件 '{img_path}' 未找到或无法读取.")
        return img
    except Exception as e:
        print(f"加载图像时出错: {e}")
        raise


def draw_bbox(img, bboxes, color=(128, 240, 128), thickness=3):
    """Draw bounding boxes on an image."""
    img_copy = img.copy()
    for i, bbox in enumerate(bboxes):
        try:
            bbox = bbox.astype(int)
            cv2.polylines(img_copy, [bbox], isClosed=True, color=color, thickness=thickness)
            # 添加文本框编号
            center = bbox.mean(axis=0).astype(int)
            cv2.putText(img_copy, str(i+1), tuple(center), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 3)
        except Exception as e:
            print(f"绘制边界框时出错: {e}")
            continue
    return img_copy


def crop_and_rotate_image(img, box):
    """Crop the image using the bounding box coordinates."""
    try:
        rect = cv2.minAreaRect(box) 
        box = cv2.boxPoints(rect).astype(np.intp)
        width = int(rect[1][0])
        height = int(rect[1][1])
        angle = rect[2]

        # 确保宽度和高度为正值
        if width <= 0 or height <= 0:
            print(f"无效的尺寸: width={width}, height={height}")
            return None

        src_pts = box.astype("float32")
        dst_pts = np.array([[0, height-1],
                            [0, 0],
                            [width-1, 0],
                            [width-1, height-1]],
                            dtype="float32")

        M = cv2.getPerspectiveTransform(src_pts, dst_pts)
        warped = cv2.warpPerspective(img, M, (width, height))
        
        # 根据角度决定是否旋转
        if angle < -45:
            rotated = cv2.rotate(warped, cv2.ROTATE_90_CLOCKWISE)
        elif angle > 45:
            rotated = cv2.rotate(warped, cv2.ROTATE_90_COUNTERCLOCKWISE)
        else:
            rotated = warped
        
        print(f"裁剪图像尺寸: 宽度={rotated.shape[1]}, 高度={rotated.shape[0]}")
        
        return rotated
    except Exception as e:
        print(f"裁剪和旋转图像时出错: {e}")
        import traceback
        traceback.print_exc()
        return None


def display_image(img, title="Image", output_path='output/predict.jpg'):
    """Display an image using Matplotlib and save it to a file."""
    try:
        # Ensure the output directory exists
        output_dir = os.path.dirname(output_path)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # Save the image to a file
        success = cv2.imwrite(output_path, img)
        if success:
            print(f"图像已保存到: {output_path}")
        else:
            print(f"保存图像失败: {output_path}")
        
        # 使用非交互式后端保存图像
        plt.figure(figsize=(15, 10))
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.title(title, fontsize=16)
        plt.axis('off')
        plt.tight_layout()
        
        # 保存matplotlib图像
        plot_path = output_path.replace('.jpg', '_plot.png')
        plt.savefig(plot_path, dpi=150, bbox_inches='tight')
        plt.close()
        print(f"matplotlib图像已保存到: {plot_path}")
        
    except Exception as e:
        print(f"显示或保存图像时出错: {e}")


def draw_text_on_image(img, texts, boxes, font_scale=1.0, color=(0, 0, 255), font_thickness=2):
    """Draw recognized texts on image based on bounding boxes."""
    for i, (text, box) in enumerate(zip(texts, boxes)):
        try:
            if not text.strip():  # 如果文本为空，显示"未识别"
                text = f"未识别_{i+1}"
            
            # Get the top-left corner of the bounding box
            top_left = box.min(axis=0).astype(int)
            
            # Draw the text above the bounding box
            cv2.putText(img, text, (top_left[0], top_left[1] - 10), 
                       cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, font_thickness)
                       
        except Exception as e:
            print(f"绘制文本时出错: {e}")
            continue


def init_args():
    parser = argparse.ArgumentParser(description='paddleocr')
    parser.add_argument('--det_model_path', 
                       default='/home/sunrise/rdk_model_zoo/rdk_model_zoo/demos/OCR/PaddleOCR/model/en_PP-OCRv3_det_640x640_nv12.bin', 
                       type=str)
    parser.add_argument('--rec_model_path', 
                       default='/home/sunrise/rdk_model_zoo/rdk_model_zoo/demos/OCR/PaddleOCR/model/en_PP-OCRv3_rec_48x320_rgb.bin', 
                       type=str)
    parser.add_argument('--image_path', 
                       default='/home/sunrise/python_project1/data/1234.jpg', 
                       type=str, help='img path for predict')
    parser.add_argument('--output_folder', 
                       default='output/predict.jpg', 
                       type=str, help='img path for output')
    args = parser.parse_args()
    return args


def main():
    detection_model = None
    recognition_model = None
    
    try:
        # Alphabet for OCR
        alphabet = """0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`abcdefghijklmnopqrstuvwxyz{|}~!"#$%&'()*+,-./  """
        converter = strLabelConverter(alphabet)

        # Load image
        print(f"正在加载图像: {args.image_path}")
        img = load_image(args.image_path)
        print(f"图像加载成功，尺寸: {img.shape}")

        # Initialize detection and recognition models
        print("正在初始化检测模型...")
        detection_model = DetectionModel(args.det_model_path)
        print("检测模型初始化成功")
        
        print("正在初始化识别模型...")
        recognition_model = rec_model(args.rec_model_path, converter)
        print("识别模型初始化成功")

        # Detect bounding boxes
        print("正在进行文本检测...")
        _, boxes_list = detection_model.predict(img, args.image_path)
        print(f"检测到 {len(boxes_list)} 个文本区域")

        if len(boxes_list) == 0:
            print("未检测到任何文本区域")
            return

        # Draw bounding boxes on the original image
        img_boxes = draw_bbox(img, boxes_list)
        display_image(img_boxes, title="检测到的文本框", output_path=args.output_folder)

        # Recognize text in each detected box
        recognized_texts = []
        
        for i, box in enumerate(boxes_list):
            print(f"\n处理第 {i + 1} 个文本框:")
            cropped_img = crop_and_rotate_image(img, box)
            
            if cropped_img is None:
                print(f"第 {i + 1} 个文本框裁剪失败")
                recognized_texts.append("")
                continue
                
            # 保存裁剪后的图像用于调试
            crop_debug_path = f'output/crop_{i+1}.jpg'
            cv2.imwrite(crop_debug_path, cropped_img)
            print(f"裁剪图像已保存到: {crop_debug_path}")
            
            raw_pred, sim_pred = recognition_model.predict_float(cropped_img)
            recognized_texts.append(sim_pred)
            print(f"原始预测: '{raw_pred}'")
            print(f"简化预测: '{sim_pred}'")

        # Create result image with text annotations
        result_img = img_boxes.copy()
        draw_text_on_image(result_img, recognized_texts, boxes_list, font_scale=2.0)

        # Display the result image
        result_path = args.output_folder.replace('.jpg', '_result.jpg')
        display_image(result_img, title="识别结果", output_path=result_path)
        
        # 输出最终结果
        print("\n=== 最终识别结果 ===")
        for i, text in enumerate(recognized_texts):
            if text.strip():
                print(f"文本框 {i+1}: '{text}'")
            else:
                print(f"文本框 {i+1}: 未识别到文本")
            
    except Exception as e:
        print(f"主程序执行时出错: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # 手动释放资源
        print("正在释放资源...")
        try:
            if detection_model is not None:
                del detection_model
            if recognition_model is not None:
                del recognition_model
        except Exception as e:
            print(f"释放资源时出错: {e}")
        print("程序执行完成")


if __name__ == '__main__':
    args = init_args()
    main()